
import math

from PIL import Image
import torch
from torch.utils.data import Dataset
import os
from resize import *
import numpy as np
from torchvision.transforms import transforms
from configure import *


tf = transforms.Compose([transforms.ToTensor()])

def on_hot(num,location):
    temp = np.zeros(num)
    temp[location] = 1
    return temp

def bbox_wh_iou(wh1, wh2):

    w1, h1 = wh1[0], wh1[1]
    w2, h2 = wh2[0], wh2[1]
    inter_area = np.min([w1, w2]) * np.min([h1, h2])
    union_area = (w1 * h1 + 1e-16) + w2 * h2 - inter_area
    return inter_area / union_area

def bbox_iou(box1, box2, x1y1x2y2=True):
    """
    Returns the IoU of two bounding boxes
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
        b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
        b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
        b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 = max(b1_x1, b2_x1)
    inter_rect_y1 = max(b1_y1, b2_y1)
    inter_rect_x2 = min(b1_x2, b2_x2)
    inter_rect_y2 = min(b1_y2, b2_y2)
    # Intersection area
    if (inter_rect_x2-inter_rect_x1)<=0:
        inter_area = 0
    else:

        inter_area  = (inter_rect_x2-inter_rect_x1+1)*(inter_rect_y2-inter_rect_y1+1)
    # Union Area
    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    return iou


class Yolo_dataset(Dataset):
    def __init__(self):
        f = open('E:/python file/yolo_advance/date.txt','r')
        self.data = f.readlines()
    def __len__(self):
        return len(self.data)
    def __getitem__(self, index):
        date = self.data[index].strip()
        date = date.split(' ')
        out_date = [ float(i) for i in date[2:]]
        out_date = np.array(out_date)
        out_date = np.split(out_date,len(out_date)//5)
        path = os.path.join('E://python file//yolo_advance//JPEGImages',date[0],date[1])
        imag_size = Image.open(path)
        scale = 416/max(imag_size.size)
        image = Yolo_resize(path)

        anchor_obj = {}
        anchor_noobj = {}
        anchor_xishu = {}
        
        for map_size ,anchor_wh in anchor.items():
            
            anchor_obj[map_size] = np.zeros((map_size,map_size,3,5+class_num),dtype=np.float32)
            anchor_noobj[map_size] = np.ones((map_size,map_size,3,1),dtype=np.float32)
            anchor_xishu[map_size] = np.zeros((map_size,map_size,3,1),dtype=np.float32)

            for target_box in out_date:
                class_,cx,cy,w,h =  target_box
            
                cx_list = [math.modf(cx*scale*x)  for x in wh_scale]
                cy_list = [math.modf(cy*scale*x) for x in wh_scale]
            
                ious = [bbox_wh_iou(x,[w,h]) for x in anchor_list ]
                best_n = np.argmax(ious,0)

                ach_feature_map = 2**((best_n)//3)*13
                ach_number = best_n % 3

                if  map_size == ach_feature_map:
                    
                    cxy_bia = int(math.log2(ach_feature_map/13))
                    
                    _w ,_h= w*scale/(anchor_wh[ach_number][0]), h*scale/(anchor_wh[ach_number][1])
                    anchor_obj[map_size][int(cy_list[cxy_bia][1]),int(cx_list[cxy_bia][1]),ach_number] = np.array([1.0,cy_list[cxy_bia][0],cx_list[cxy_bia][0],np.log(_w+1e-16),np.log(_h+1e-16),*on_hot(class_num,int(class_))])
                    anchor_noobj[map_size][int(cy_list[cxy_bia][1]),int(cx_list[cxy_bia][1]),ach_number] = 0.0
                    anchor_xishu[map_size][int(cy_list[cxy_bia][1]),int(cx_list[cxy_bia][1]),ach_number] = (2 - w*h*(scale/416)**2)

                    for i in range(map_size):
                        for j in range(map_size):
                            for k in range(3):
                                if bbox_iou([j,i,anchor_wh[k][0]*map_size/416,anchor_wh[k][1]*map_size/416],[cx*scale*map_size/416,cy*scale*map_size/416,w*map_size/416,h*map_size/416],False)>0.5:
                                    anchor_noobj[map_size][i,j,k] = 0.0

        return tf(image),anchor_obj[13],anchor_obj[26],anchor_obj[52],anchor_noobj[13],anchor_noobj[26],anchor_noobj[52],anchor_xishu[13],anchor_xishu[26],anchor_xishu[52]

if __name__ =='__main__':
   yolo = Yolo_dataset()
   mask = yolo[0][1][...,0] == 1
   mask1 = yolo[0][2][...,0] == 1
   for i in range(98):

        print(i)
        mask2 = yolo[i][4][...,0] == 0
#    print(yolo[0][1][mask])
#    print(yolo[0][2][mask1])
#    print(yolo[0][3][mask2])
        print(yolo[i][4][mask2])
#    print(yolo[0][5].shape)
#    print(yolo[0][6].shape)

   